from datetime import datetime
from typing import Optional, List
from apps.models import Student_trainings, Student, Coach, Training_Programs, Student_Training_Status, Rating, User
from tortoise.transactions import in_transaction
from tortoise.exceptions import DoesNotExist


async def get_student_trainings_by_user_id(user_id: int):
    try:
        # 先根据 user_id 找到对应的学生或教练
        student = await Student.get_or_none(user_id=user_id)
        coach = await Coach.get_or_none(user_id=user_id)

        # 如果既不是学生也不是教练
        if not student and not coach:
            raise ValueError("User is neither a student nor a coach")

        # 根据用户类型查询训练记录
        if student:
            # 学生查询自己的训练记录
            trainings = await Student_trainings.filter(student_id_id=student.student_id).all().prefetch_related(
                'student_id__user_id', 'coach_id__user_id', 'program_id'
            )
            user_type = "student"
            user_obj = student
        else:
            # 教练查询自己负责的训练记录
            trainings = await Student_trainings.filter(coach_id_id=coach.coach_id).all().prefetch_related(
                'student_id__user_id', 'coach_id__user_id', 'program_id'
            )
            user_type = "coach"
            user_obj = coach

        if not trainings:
            raise ValueError(f"No training records found for this {user_type}")

        # 格式化训练记录数据
        training_list = []
        for training in trainings:
            try:
                # 获取关联的教练信息
                coach_user = None
                if training.coach_id:
                    coach_user = await User.get_or_none(user_id=training.coach_id.user_id_id)

                # 获取关联的学生信息
                student_user = None
                if training.student_id:
                    student_user = await User.get_or_none(user_id=training.student_id.user_id_id)

                # 获取关联的项目信息
                program = await training.program_id

                training_list.append({
                    "training_id": training.training_id,
                    "student_id": training.student_id_id,
                    "student_name": student_user.username if student_user else "Unknown",
                    "coach_id": training.coach_id_id,
                    "coach_name": coach_user.username if coach_user else "Unknown",
                    "program_id": training.program_id_id,
                    "program_name": program.name if program else "Unknown",
                    "start_time": training.start_time.isoformat() if training.start_time else None,
                    "end_time": training.end_time.isoformat() if training.end_time else None,
                    "status": training.status.value if hasattr(training.status, 'value') else training.status,
                    "feedback": training.feedback,
                    "rating": training.rating.value if hasattr(training.rating, 'value') else training.rating,
                    "content": training.content,
                    "response": training.response,
                    "create_at": training.create_at.isoformat() if training.create_at else None
                })
            except Exception as e:
                print(f"Error processing training {training.training_id}: {str(e)}")
                continue

        return {
            "data": training_list,
            "total": len(training_list),
            "user_id": user_id,
            "user_type": user_type,
            f"{user_type}_id": user_obj.student_id if user_type == "student" else user_obj.coach_id
        }
    except Exception as e:
        print(f"Error in get_student_trainings_by_user_id: {str(e)}")
        raise


async def get_all_student_trainings(
        student_id: Optional[int] = None,
        coach_id: Optional[int] = None,
        program_id: Optional[int] = None,
        status: Optional[str] = None,
        page: int = 1,
        page_size: int = 10
):
    # 构建查询条件
    query = Student_trainings.all()

    if student_id is not None:
        query = query.filter(student_id_id=student_id)

    if coach_id is not None:
        query = query.filter(coach_id_id=coach_id)

    if program_id is not None:
        query = query.filter(program_id_id=program_id)

    if status is not None:
        query = query.filter(status=status)

    # 计算偏移量
    offset = (page - 1) * page_size

    # 获取总数
    total = await query.count()

    # 查询当前页的数据
    trainings = await query.offset(offset).limit(page_size).prefetch_related(
        'student_id__user_id',
        'coach_id__user_id',
        'program_id'
    )

    # 处理外键字段名
    formatted_trainings = []
    for training in trainings:
        # 获取关联信息
        student_user = await User.get_or_none(user_id=training.student_id.user_id_id) if training.student_id else None
        coach_user = await User.get_or_none(user_id=training.coach_id.user_id_id) if training.coach_id else None
        program = await training.program_id

        formatted_trainings.append({
            "training_id": training.training_id,
            "student_id": training.student_id_id,
            "student_name": student_user.username if student_user else "Unknown",
            "coach_id": training.coach_id_id,
            "coach_name": coach_user.username if coach_user else "Unknown",
            "program_id": training.program_id_id,
            "program_name": program.name if program else "Unknown",
            "start_time": training.start_time.isoformat() if training.start_time else None,
            "end_time": training.end_time.isoformat() if training.end_time else None,
            "status": training.status.value if hasattr(training.status, 'value') else training.status,
            "feedback": training.feedback,
            "rating": training.rating.value if hasattr(training.rating, 'value') else training.rating,
            "content": training.content,
            "response": training.response,
            "create_at": training.create_at.isoformat() if training.create_at else None
        })

    return total, formatted_trainings


async def get_student_training_by_id(training_id: int):
    try:
        training = await Student_trainings.get(training_id=training_id)

        # 获取关联信息
        student_user = await User.get_or_none(user_id=training.student_id.user_id_id) if training.student_id else None
        coach_user = await User.get_or_none(user_id=training.coach_id.user_id_id) if training.coach_id else None
        program = await training.program_id

        return {
            "training_id": training.training_id,
            "student_id": training.student_id_id,
            "student_name": student_user.username if student_user else "Unknown",
            "coach_id": training.coach_id_id,
            "coach_name": coach_user.username if coach_user else "Unknown",
            "program_id": training.program_id_id,
            "program_name": program.name if program else "Unknown",
            "start_time": training.start_time.isoformat() if training.start_time else None,
            "end_time": training.end_time.isoformat() if training.end_time else None,
            "status": training.status.value if hasattr(training.status, 'value') else training.status,
            "feedback": training.feedback,
            "rating": training.rating.value if hasattr(training.rating, 'value') else training.rating,
            "content": training.content,
            "response": training.response,
            "create_at": training.create_at.isoformat() if training.create_at else None
        }
    except DoesNotExist:
        return None


async def create_student_training(data):
    # 检查学生是否存在
    student = await Student.get_or_none(student_id=data.student_id)
    if not student:
        raise ValueError("Student not found")

    # 检查教练是否存在
    coach = await Coach.get_or_none(coach_id=data.coach_id)
    if not coach:
        raise ValueError("Coach not found")

    # 检查训练项目是否存在
    program = await Training_Programs.get_or_none(program_id=data.program_id)
    if not program:
        raise ValueError("Training program not found")

    try:
        async with in_transaction() as connection:
            # 创建学生训练记录
            training = await Student_trainings.create(
                student_id_id=student.student_id,
                coach_id_id=coach.coach_id,
                program_id_id=program.program_id,
                start_time=data.start_time,
                end_time=data.end_time,
                status=data.status,
                feedback=data.feedback,
                rating=data.rating,
                content=data.content,
                response=data.response,
                using_db=connection
            )

        return training

    except Exception as e:
        raise Exception(f"Error creating student training record: {str(e)}")


async def update_student_training(training_id: int, data):
    # 查询训练记录
    try:
        training = await Student_trainings.get(training_id=training_id)
    except DoesNotExist:
        return False

    update_data = data.dict(exclude_unset=True)  # 只获取提供的字段

    # 更新训练记录表
    update_fields = {}
    for field in ['start_time', 'end_time', 'status', 'feedback', 'rating', 'content', 'response']:
        if field in update_data:
            update_fields[field] = update_data[field]

    if update_fields:
        await training.update_from_dict(update_fields).save()

    # 更新外键关系（如果提供了）
    if 'student_id' in update_data:
        student = await Student.get_or_none(student_id=update_data['student_id'])
        if not student:
            raise ValueError("Student not found")
        training.student_id_id = student.student_id

    if 'coach_id' in update_data:
        coach = await Coach.get_or_none(coach_id=update_data['coach_id'])
        if not coach:
            raise ValueError("Coach not found")
        training.coach_id_id = coach.coach_id

    if 'program_id' in update_data:
        program = await Training_Programs.get_or_none(program_id=update_data['program_id'])
        if not program:
            raise ValueError("Training program not found")
        training.program_id_id = program.program_id

    await training.save()
    return True


async def delete_student_training(training_id: int):
    try:
        training = await Student_trainings.get(training_id=training_id)
        await training.delete()
        return True
    except DoesNotExist:
        return False
